package com.kwdb.sdk.batch;

import com.kwdb.sdk.config.ClientConfig;
import com.kwdb.sdk.data.DataPoint;
import com.kwdb.sdk.grpc.KWDBGrpcClient;
import com.kwdb.sdk.exception.KWDBException;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 批量写入器 - 高性能批量处理
 * 
 * 功能特性：
 * - 内存缓冲队列
 * - 定时/定量触发
 * - 自动重试机制
 * - 背压控制
 */
public class BatchWriter {
    
    private final ClientConfig config;
    private final KWDBGrpcClient grpcClient;
    private final BlockingQueue<DataPoint> dataQueue;
    private final ScheduledExecutorService scheduler;
    private final AtomicBoolean running = new AtomicBoolean(false);
    private final AtomicLong totalWritten = new AtomicLong(0);
    private final AtomicLong totalErrors = new AtomicLong(0);
    
    public BatchWriter(ClientConfig config, KWDBGrpcClient grpcClient) {
        this.config = config;
        this.grpcClient = grpcClient;
        this.dataQueue = new LinkedBlockingQueue<>();
        this.scheduler = Executors.newScheduledThreadPool(1);
    }
    
    /**
     * 启动批量写入器
     */
    public void start() {
        if (running.compareAndSet(false, true)) {
            // 启动定时批量提交任务
            scheduler.scheduleAtFixedRate(this::flushBatch, 
                config.getBatchTimeout(), 
                config.getBatchTimeout(), 
                TimeUnit.MILLISECONDS);
        }
    }
    
    /**
     * 停止批量写入器
     */
    public void stop() {
        if (running.compareAndSet(true, false)) {
            // 提交剩余数据
            flushBatch();
            scheduler.shutdown();
            try {
                if (!scheduler.awaitTermination(5, TimeUnit.SECONDS)) {
                    scheduler.shutdownNow();
                }
            } catch (InterruptedException e) {
                scheduler.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }
    }
    
    /**
     * 添加数据点
     * @param dataPoint 数据点
     */
    public void addDataPoint(DataPoint dataPoint) {
        if (running.get()) {
            try {
                dataQueue.offer(dataPoint, 100, TimeUnit.MILLISECONDS);
                
                // 如果队列大小达到批量大小，立即提交
                if (dataQueue.size() >= config.getBatchSize()) {
                    flushBatch();
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                throw new RuntimeException("添加数据点被中断", e);
            }
        }
    }
    
    /**
     * 添加批量数据点
     * @param dataPoints 数据点列表
     */
    public void addBatch(List<DataPoint> dataPoints) {
        if (running.get()) {
            try {
                for (DataPoint dataPoint : dataPoints) {
                    dataQueue.offer(dataPoint, 100, TimeUnit.MILLISECONDS);
                }
                
                // 如果队列大小达到批量大小，立即提交
                if (dataQueue.size() >= config.getBatchSize()) {
                    flushBatch();
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                throw new RuntimeException("添加批量数据点被中断", e);
            }
        }
    }
    
    /**
     * 强制提交当前队列中的所有数据
     */
    public void flushBatch() {
        if (!running.get() || dataQueue.isEmpty()) {
            return;
        }
        
        List<DataPoint> batch = new ArrayList<>();
        dataQueue.drainTo(batch, config.getBatchSize());
        
        if (batch.isEmpty()) {
            return;
        }
        
        // 重试机制
        int retries = 0;
        Exception lastException = null;
        
        while (retries < config.getMaxRetries()) {
            try {
                grpcClient.writeBatch(batch);
                totalWritten.addAndGet(batch.size());
                return; // 成功，退出重试循环
                
            } catch (Exception e) {
                lastException = e;
                retries++;
                
                if (retries < config.getMaxRetries()) {
                    try {
                        Thread.sleep(config.getRetryInterval());
                    } catch (InterruptedException ie) {
                        Thread.currentThread().interrupt();
                        throw new RuntimeException("重试被中断", ie);
                    }
                }
            }
        }
        
        // 所有重试都失败了
        totalErrors.addAndGet(batch.size());
        System.err.println("批量写入失败，已重试 " + config.getMaxRetries() + " 次: " + 
                          (lastException != null ? lastException.getMessage() : "未知错误"));
        
        // 可以选择将失败的数据重新放回队列或丢弃
        // 这里选择丢弃，避免无限重试
    }
    
    /**
     * 获取队列大小
     * @return 队列中的数据点数量
     */
    public int getQueueSize() {
        return dataQueue.size();
    }
    
    /**
     * 检查是否正在运行
     * @return 是否正在运行
     */
    public boolean isRunning() {
        return running.get();
    }
    
    /**
     * 获取统计信息
     * @return 统计信息字符串
     */
    public String getStats() {
        return String.format("队列大小: %d, 已写入: %d, 错误数: %d", 
                           getQueueSize(), 
                           totalWritten.get(), 
                           totalErrors.get());
    }
    
    /**
     * 重置统计信息
     */
    public void resetStats() {
        totalWritten.set(0);
        totalErrors.set(0);
    }
}
